use std::{
    env,
    ffi::CString,
    fs,
    fs::File,
    io::{Read, Seek, SeekFrom},
    num::NonZero,
    os::unix::ffi::OsStrExt,
    path::PathBuf,
    time::Duration,
};

use libafl::{
    corpus::{InMemoryCorpus, OnDiskCorpus},
    events::SimpleEventManager,
    executors::{
        command::{CommandConfigurator, PTraceCommandConfigurator},
        hooks::intel_pt::{IntelPT, IntelPTHook, PtImage, PAGE_SIZE},
    },
    feedbacks::{CrashFeedback, MaxMapFeedback},
    fuzzer::{Fuzzer, StdFuzzer},
    generators::RandPrintablesGenerator,
    monitors::SimpleMonitor,
    mutators::{havoc_mutations::havoc_mutations, scheduled::HavocScheduledMutator},
    observers::ConstMapObserver,
    schedulers::QueueScheduler,
    stages::mutational::StdMutationalStage,
    state::StdState,
};
use libafl_bolts::{core_affinity, nonnull_raw_mut, rands::StdRand, tuples::tuple_list, Error};
use object::{elf::PF_X, Object, ObjectSegment, SegmentFlags};

// Coverage map
const MAP_SIZE: usize = 4096;
static mut MAP: [u8; MAP_SIZE] = [0; MAP_SIZE];
static mut MAP_PTR: *mut u8 = &raw mut MAP as _;

pub fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Let's set the default logging level to `warn`
    if env::var("RUST_LOG").is_err() {
        env::set_var("RUST_LOG", "warn")
    }
    // Enable logging
    env_logger::init();

    // path of the program we want to fuzz
    let target_path = PathBuf::from(env::args().next().unwrap())
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .join("target_program");
    let target_binary = fs::read(&target_path)?;
    let target_parsed = object::File::parse(&*target_binary)?;

    // Get the executable segment from the target program.
    //
    // Note: this simple example target has just one range of executable memory, in real world binaries
    // there are likely multiple ranges.
    let executable_segments = target_parsed.segments().filter(|s| match s.flags() {
        SegmentFlags::Elf { p_flags } => (p_flags & PF_X) > 0,
        _ => panic!("target binary is not an ELF file."),
    });
    let executable_segment = executable_segments
        .into_iter()
        .next()
        .ok_or_else(|| Error::illegal_argument("No executable segment found in target program"))?;
    log::debug!(
        "Executable segment: {executable_segment:x?} at binary file offset {:x?}",
        executable_segment.file_range()
    );

    // We'll run the target on cpu (aka core) 0
    let cpu = core_affinity::get_core_ids()?[0];
    log::debug!("Using core {} for fuzzing", cpu.0);

    // Create an observation channel using the map
    let observer = unsafe { ConstMapObserver::from_mut_ptr("signals", nonnull_raw_mut!(MAP)) };

    // Feedback to rate the interestingness of an input
    let mut feedback = MaxMapFeedback::new(&observer);

    // A feedback to choose if an input is a solution or not
    let mut objective = CrashFeedback::new();

    // create a State from scratch
    let mut state = StdState::new(
        // RNG
        StdRand::new(),
        // Corpus that will be evolved, we keep it in memory for performance
        InMemoryCorpus::new(),
        // Corpus in which we store solutions (crashes in this example),
        // on disk so the user can get them after stopping the fuzzer
        OnDiskCorpus::new(PathBuf::from("./crashes")).unwrap(),
        // States of the feedbacks.
        // The feedbacks can report the data that should persist in the State.
        &mut feedback,
        // Same for objective feedbacks
        &mut objective,
    )?;

    // The Monitor trait define how the fuzzer stats are displayed to the user
    let mon = SimpleMonitor::new(|s| println!("{s}"));

    // The event manager handle the various events generated during the fuzzing loop
    // such as the notification of the addition of a new item to the corpus
    let mut mgr = SimpleEventManager::new(mon);

    // A queue policy to get testcases from the corpus
    let scheduler = QueueScheduler::new();

    // A fuzzer with feedbacks and a corpus scheduler
    let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);

    // The target is a ET_DYN elf, it will be relocated by the loader with this offset.
    // see https://github.com/torvalds/linux/blob/c1e939a21eb111a6d6067b38e8e04b8809b64c4e/arch/x86/include/asm/elf.h#L234C1-L239C38
    const DEFAULT_MAP_WINDOW: u64 = (1 << 47) - PAGE_SIZE as u64;
    const ELF_ET_DYN_BASE: u64 = (DEFAULT_MAP_WINDOW / 3 * 2) & !(PAGE_SIZE as u64 - 1);

    // Set the instruction pointer (IP) filter and memory image of our target.
    let actual_virtual_address = executable_segment.address() + ELF_ET_DYN_BASE;
    let filters = vec![actual_virtual_address..=actual_virtual_address + executable_segment.size()];

    let mut file = File::open(&target_path)?;
    let mut data = vec![0; executable_segment.size() as usize];
    file.seek(SeekFrom::Start(executable_segment.file_range().0))?;
    file.read_exact(&mut data)?;
    let images = vec![PtImage::new(data, actual_virtual_address)];

    let intel_pt = IntelPT::builder()
        .cpu(cpu.0)
        .inherit(true)
        .ip_filters(filters)
        .images(images)
        .build()?;

    let hook = unsafe { IntelPTHook::builder().map_ptr(MAP_PTR).map_len(MAP_SIZE) }
        .intel_pt(intel_pt)
        .build();

    let target_cstring = CString::from(
        target_path
            .as_os_str()
            .as_bytes()
            .iter()
            .map(|&b| NonZero::new(b).unwrap())
            .collect::<Vec<_>>(),
    );

    let command_configurator = PTraceCommandConfigurator::builder()
        .path(target_cstring)
        .cpu(cpu)
        .timeout(Duration::from_secs(2))
        .build();
    let mut executor = command_configurator.into_executor_with_hooks(
        tuple_list!(observer),
        tuple_list!(hook),
        None,
        None,
    );

    // Generator of printable bytearrays of max size 32
    let mut generator = RandPrintablesGenerator::new(NonZero::new(32).unwrap());

    // Generate 8 initial inputs
    state
        .generate_initial_inputs(&mut fuzzer, &mut executor, &mut generator, &mut mgr, 8)
        .expect("Failed to generate the initial corpus");

    // Setup a mutational stage with a basic bytes mutator
    let mutator = HavocScheduledMutator::new(havoc_mutations());
    let mut stages = tuple_list!(StdMutationalStage::new(mutator));

    fuzzer
        .fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)
        .expect("Error in the fuzzing loop");

    Ok(())
}
